# TODO: 使用TensorFlow实现训练并将训练结果保存到mnist.model
# DATE: 2022/3/24
# AUTHOR: Cheng Ze WUST

import tensorflow as tf
import os
os.environ['(TF_CPP_MIN)LOG_LEVEL']='2'
minist=tf.keras.datasets.mnist

def create_model():
    #直接实现神经网络模型
    model=tf.keras.models.Sequential([
        # 设置第一个隐藏层，有50个结点，使用sigmoid激活函数，输入有784个
        tf.keras.layers.Dense(50,activation='sigmoid',input_shape=(784,)),
        # 设置第二个隐藏层，有100个结点，使用sigmoid激活函数，输入有由第一隐藏层传来的结果50个
        tf.keras.layers.Dense(100, activation='sigmoid', input_shape=(50,)),
        # 为了防止过拟合，设置丢弃的比例20%（随机将输入单元设置为0）
        tf.keras.layers.Dropout(0.2),
        tf.keras.layers.Dense(10, activation='softmax')
    ])
    model.compile(optimizer='adam',loss='sparse_categorical_crossentropy',metrics=['accuracy'])
    return model

model=create_model()
(train_img,train_lables),(test_img,test_lables)=tf.keras.datasets.mnist.load_data()
train_lables=train_lables[:10000]   #训练数据取前10000条
test_lables=test_lables[:10000]
train_img=train_img[:10000].reshape(-1,28*28)/255.0
test_img=test_img[:10000].reshape(-1,28*28)/255.0

model.fit(train_img,train_lables,epochs=50) #迭代次数20，由于选择的损失函数故迭代次数不用太大

model.save('mnist.model')




